import torch
from torchvision import transforms
from PIL import Image
from .models import CNNMnist, CNNDropoutMnist
import os


class MnistService:
    def __init__(self):
        self.device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
        self.model = CNNMnist().to(self.device)
        self.model.load_state_dict(torch.load('mnist/ml/mnist_cnn.pth'))
        self.model.eval()

        self.transform = transforms.Compose([
            transforms.Resize((28, 28)),
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])
    def predict(self, image_path):
        image = Image.open(image_path).convert('L')
        img_tensor = self.transform(image).unsqueeze(0).to(self.device)
        with torch.no_grad():
            output = self.model(img_tensor)
            prediction = output.argmax(dim=1).item()
            confidence = torch.exp(output.max()).item()
            print("手写数字图片识别结果:", prediction, confidence)
        return prediction, confidence


## 加入Dropout解决过拟合问题
class MnistDropoutService:
    def __init__(self):
        self.device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
        self.model = CNNDropoutMnist().to(self.device)
        self.model.load_state_dict(torch.load('mnist/ml/mnist_cnn_dropout.pth'))
        self.model.eval()

        self.transform = transforms.Compose([
            transforms.Resize((28, 28)),
            transforms.Grayscale(), # 添加灰度化,避免前端上传非黑白图片导致识别失败
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])

    def predict(self, image_path):
        image = Image.open(image_path).convert('L')
        img_tensor = self.transform(image).unsqueeze(0).to(self.device)
        with torch.no_grad():
            output = self.model(img_tensor)
        prob = torch.nn.functional.softmax(output, dim=1)
        prediction =prob.argmax().item()
        confidence = prob.max().item()
        print("手写数字图片识别结果:", prediction, confidence)
        return prediction, confidence